{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Named Entity Recognition for Security Intelligence\n",
    "\n",
    "A security-focused NER pipeline using Fenic's semantic extraction capabilities to identify and analyze threats, vulnerabilities, and indicators of compromise from unstructured security reports.\n",
    "\n",
    "This pipeline demonstrates automated security entity extraction and risk assessment:\n",
    "\n",
    "- Zero-shot entity extraction (CVEs, IPs, domains, hashes)\n",
    "- Enhanced extraction with threat intelligence context\n",
    "- Document chunking for comprehensive analysis\n",
    "- Risk prioritization and actionable intelligence"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup the session and manage imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import fenic as fc\n",
    "import re\n",
    "from pydantic import BaseModel, Field\n",
    "\n",
    "config = fc.SessionConfig(\n",
    "    app_name=\"security_vulnerability_ner\",\n",
    "    semantic=fc.SemanticConfig(\n",
    "        language_models={\n",
    "            \"mini\" : fc.OpenAIModelConfig(model_name=\"gpt-4o-mini\", rpm=500, tpm=200_000)\n",
    "        }\n",
    "    )\n",
    ")\n",
    "\n",
    "session = fc.Session.get_or_create(config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sample Vulnerability Reports Dataset\n",
    "\n",
    "This cell defines a sample dataset of security vulnerability reports. Each report includes details such as report ID, source, title, and content describing real-world security incidents, vulnerabilities, and threat intelligence. The dataset is used to demonstrate entity extraction and analysis in the NER pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sample vulnerability reports data\n",
    "vulnerability_reports_data = [\n",
    "        {\n",
    "            \"report_id\": \"CVE-2024-001\",\n",
    "            \"source\": \"CVE Database\",\n",
    "            \"title\": \"Critical OpenSSL Buffer Overflow\",\n",
    "            \"content\": \"CVE-2024-3094: Buffer overflow in OpenSSL 3.0.0-3.0.12. Affects Ubuntu 22.04, RHEL 8. CVSS 9.8. IOCs: evil-domain.com, 10.0.0.50:443\"\n",
    "        },\n",
    "        {\n",
    "            \"report_id\": \"THREAT-2024-002\",\n",
    "            \"source\": \"Threat Intelligence\",\n",
    "            \"title\": \"APT29 Campaign Targeting Financial Sector\",\n",
    "            \"content\": \"APT29 targeting banks. Exploits CVE-2024-1234, CVE-2024-5678. Malware: SUNBURST 2.0. C2: c2-server.badguys.net (185.159.158.1)\"\n",
    "        },\n",
    "        {\n",
    "            \"report_id\": \"SEC-ADV-2024-003\",\n",
    "            \"source\": \"Security Advisory\",\n",
    "            \"title\": \"Zero-Day in Popular WordPress Plugin\",\n",
    "            \"content\": \"CVE-2024-9999: SQL injection in WP Super Cache 1.0.0-1.7.8. CVSS 8.5. Patch in 1.7.9. Related: CVE-2024-9998, CVE-2024-9997\"\n",
    "        },\n",
    "        {\n",
    "            \"report_id\": \"INC-2024-004\",\n",
    "            \"source\": \"Incident Report\",\n",
    "            \"title\": \"Ransomware Attack on Healthcare Provider\",\n",
    "            \"content\": \"LockBit 3.0 ransomware via CVE-2024-4444. Used Mimikatz 2.2.0. C2: 45.142.214.99:8443. Affected Windows Server 2016, 2019\"\n",
    "        },\n",
    "        {\n",
    "            \"report_id\": \"VULN-2024-005\",\n",
    "            \"source\": \"Bug Bounty Report\",\n",
    "            \"title\": \"Authentication Bypass in Enterprise SaaS Platform\",\n",
    "            \"content\": \"Auth bypass in AuthProvider 2.5.1 at login.platform.com. JWT alg:none vulnerability. Fixed in 2.5.2. Affects /api/v2/admin/*\"\n",
    "        }\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create DataFrame\n",
    "reports_df = session.create_dataframe(vulnerability_reports_data)\n",
    "\n",
    "print(\"🔒 Security Vulnerability NER Pipeline\")\n",
    "print(\"=\" * 70)\n",
    "print(f\"Processing {reports_df.count()} vulnerability reports\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stage 1: Basic Zero-Shot Entity Extraction\n",
    "\n",
    "This cell defines a basic Named Entity Recognition (NER) schema for extracting key security-related entities—such as CVE IDs, software packages, IP addresses, domains, and file hashes—from vulnerability report content using Fenic's zero-shot extraction. It applies the schema to the dataset, extracts entities, and displays sample results for review."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Stage 1: Basic NER with zero-shot extraction\n",
    "print(\"🔍 Stage 1: Zero-shot entity extraction...\")\n",
    "\n",
    "# Define basic NER schema for security entities\n",
    "basic_ner_schema = fc.ExtractSchema([\n",
    "    fc.ExtractSchemaField(\n",
    "        name=\"cve_ids\",\n",
    "        data_type=fc.ExtractSchemaList(element_type=fc.StringType),\n",
    "        description=\"CVE identifiers in format CVE-YYYY-NNNNN\"\n",
    "    ),\n",
    "    fc.ExtractSchemaField(\n",
    "        name=\"software_packages\",\n",
    "        data_type=fc.ExtractSchemaList(element_type=fc.StringType),\n",
    "        description=\"Software names and versions mentioned\"\n",
    "    ),\n",
    "    fc.ExtractSchemaField(\n",
    "        name=\"ip_addresses\",\n",
    "        data_type=fc.ExtractSchemaList(element_type=fc.StringType),\n",
    "        description=\"IP addresses (IPv4 or IPv6)\"\n",
    "    ),\n",
    "    fc.ExtractSchemaField(\n",
    "        name=\"domains\",\n",
    "        data_type=fc.ExtractSchemaList(element_type=fc.StringType),\n",
    "        description=\"Domain names and URLs\"\n",
    "    ),\n",
    "    fc.ExtractSchemaField(\n",
    "        name=\"file_hashes\",\n",
    "        data_type=fc.ExtractSchemaList(element_type=fc.StringType),\n",
    "        description=\"File hashes (MD5, SHA1, SHA256)\"\n",
    "    )\n",
    "])\n",
    "\n",
    "# Apply basic extraction\n",
    "basic_extraction_df = reports_df.select(\n",
    "    \"report_id\",\n",
    "    \"source\",\n",
    "    \"title\",\n",
    "    fc.semantic.extract(\"content\", basic_ner_schema).alias(\"basic_entities\")\n",
    ")\n",
    "\n",
    "# Display sample results\n",
    "print(\"Sample basic extraction results:\")\n",
    "basic_readable = basic_extraction_df.select(\n",
    "    \"report_id\",\n",
    "    basic_extraction_df.basic_entities.cve_ids.alias(\"cve_ids\"),\n",
    "    basic_extraction_df.basic_entities.software_packages.alias(\"software_packages\")\n",
    ")\n",
    "basic_readable.show(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stage 2: Enhanced Domain-Specific Entity Extraction\n",
    "\n",
    "This cell defines an advanced NER schema that expands on the basic extraction by including additional security-specific entities such as attack vectors, threat actors, CVSS scores, MITRE techniques, and affected systems. It preprocesses the report content for consistency, applies the enhanced schema to extract richer security intelligence, and displays sample results for key entities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Stage 2: Enhanced extraction with domain-specific schema\n",
    "print(\"\\n🧠 Stage 2: Enhanced domain-specific extraction...\")\n",
    "\n",
    "# Define enhanced schema with security-specific entities\n",
    "enhanced_ner_schema = fc.ExtractSchema([\n",
    "    # Include all basic entities\n",
    "    fc.ExtractSchemaField(\n",
    "        name=\"cve_ids\",\n",
    "        data_type=fc.ExtractSchemaList(element_type=fc.StringType),\n",
    "        description=\"CVE identifiers in format CVE-YYYY-NNNNN\"\n",
    "    ),\n",
    "    fc.ExtractSchemaField(\n",
    "        name=\"software_packages\",\n",
    "        data_type=fc.ExtractSchemaList(element_type=fc.StringType),\n",
    "        description=\"Software names with specific version numbers\"\n",
    "    ),\n",
    "    fc.ExtractSchemaField(\n",
    "        name=\"ip_addresses\",\n",
    "        data_type=fc.ExtractSchemaList(element_type=fc.StringType),\n",
    "        description=\"IP addresses (IPv4 or IPv6)\"\n",
    "    ),\n",
    "    fc.ExtractSchemaField(\n",
    "        name=\"domains\",\n",
    "        data_type=fc.ExtractSchemaList(element_type=fc.StringType),\n",
    "        description=\"Domain names, subdomains, and URLs\"\n",
    "    ),\n",
    "    fc.ExtractSchemaField(\n",
    "        name=\"file_hashes\",\n",
    "        data_type=fc.ExtractSchemaList(element_type=fc.StringType),\n",
    "        description=\"File hashes with hash type prefix (MD5:, SHA1:, SHA256:)\"\n",
    "    ),\n",
    "    # Additional security-specific entities\n",
    "    fc.ExtractSchemaField(\n",
    "        name=\"attack_vectors\",\n",
    "        data_type=fc.ExtractSchemaList(element_type=fc.StringType),\n",
    "        description=\"Attack methods like buffer overflow, SQL injection, phishing\"\n",
    "    ),\n",
    "    fc.ExtractSchemaField(\n",
    "        name=\"threat_actors\",\n",
    "        data_type=fc.ExtractSchemaList(element_type=fc.StringType),\n",
    "        description=\"Threat actor names, APT groups, ransomware families\"\n",
    "    ),\n",
    "    fc.ExtractSchemaField(\n",
    "        name=\"cvss_scores\",\n",
    "        data_type=fc.ExtractSchemaList(element_type=fc.StringType),\n",
    "        description=\"CVSS scores and severity ratings\"\n",
    "    ),\n",
    "    fc.ExtractSchemaField(\n",
    "        name=\"mitre_techniques\",\n",
    "        data_type=fc.ExtractSchemaList(element_type=fc.StringType),\n",
    "        description=\"MITRE ATT&CK technique IDs (TXXXX format)\"\n",
    "    ),\n",
    "    fc.ExtractSchemaField(\n",
    "        name=\"affected_systems\",\n",
    "        data_type=fc.ExtractSchemaList(element_type=fc.StringType),\n",
    "        description=\"Operating systems, platforms, or infrastructure affected\"\n",
    "    )\n",
    "])\n",
    "\n",
    "# Preprocess content for better extraction\n",
    "@fc.udf(return_type=fc.StringType)\n",
    "def preprocess_udf(content):\n",
    "    # Standardize CVE format\n",
    "    content = re.sub(r'CVE\\s*-\\s*(\\d{4})\\s*-\\s*(\\d+)', r'CVE-\\1-\\2', content)\n",
    "    # Normalize version ranges\n",
    "    content = re.sub(r'(\\d+\\.\\d+\\.\\d+)\\s+through\\s+(\\d+\\.\\d+\\.\\d+)', r'\\1 to \\2', content)\n",
    "    # Clean up extra whitespace\n",
    "    content = ' '.join(content.split())\n",
    "    return content\n",
    "\n",
    "# Apply preprocessing and enhanced extraction\n",
    "enhanced_df = reports_df.select(\n",
    "    \"report_id\",\n",
    "    \"source\",\n",
    "    \"title\",\n",
    "    \"content\",\n",
    "    preprocess_udf(\"content\").alias(\"processed_content\")\n",
    ").select(\n",
    "    \"report_id\",\n",
    "    \"source\",\n",
    "    \"title\",\n",
    "    \"content\",\n",
    "    fc.semantic.extract(\"processed_content\", enhanced_ner_schema).alias(\"entities\")\n",
    ")\n",
    "\n",
    "print(\"Enhanced extraction with security-specific entities:\")\n",
    "enhanced_readable = enhanced_df.select(\n",
    "    \"report_id\",\n",
    "    enhanced_df.entities.threat_actors.alias(\"threat_actors\"),\n",
    "    enhanced_df.entities.attack_vectors.alias(\"attack_vectors\"),\n",
    "    enhanced_df.entities.cvss_scores.alias(\"cvss_scores\")\n",
    ")\n",
    "enhanced_readable.show(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stage 3: Chunking and Processing Long Documents\n",
    "\n",
    "This cell demonstrates how to handle long vulnerability reports by splitting (chunking) their content into smaller, overlapping segments for more effective entity extraction. It identifies which reports need chunking, applies recursive word chunking to long documents, extracts entities from each chunk using the enhanced schema, aggregates the results, and displays a sample of the aggregated entities per report."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Stage 3: Process long documents with chunking\n",
    "print(\"\\n📄 Stage 3: Chunking and processing long documents...\")\n",
    "\n",
    "# Add content length for chunking decisions\n",
    "reports_with_length = enhanced_df.select(\n",
    "    \"*\",\n",
    "    fc.text.length(fc.col(\"content\")).alias(\"content_length\")\n",
    ")\n",
    "\n",
    "# Identify documents needing chunking (>80 characters for demo)\n",
    "long_reports = reports_with_length.filter(fc.col(\"content_length\") > 80)\n",
    "short_reports = reports_with_length.filter(fc.col(\"content_length\") <= 80)\n",
    "\n",
    "print(f\"Documents requiring chunking: {long_reports.count()}\")\n",
    "print(f\"Documents processed whole: {short_reports.count()}\")\n",
    "\n",
    "# Apply chunking to long documents\n",
    "chunked_df = long_reports.select(\n",
    "    \"report_id\",\n",
    "    \"content\",\n",
    "    fc.text.recursive_word_chunk(\n",
    "        fc.col(\"content\"),\n",
    "        chunk_size=50,\n",
    "        chunk_overlap_percentage=15\n",
    "    ).alias(\"chunks\")\n",
    ").explode(\"chunks\").select(\n",
    "    \"report_id\",\n",
    "    fc.col(\"chunks\").alias(\"chunk\")\n",
    ")\n",
    "\n",
    "# Extract entities from each chunk\n",
    "chunk_entities_df = chunked_df.select(\n",
    "    \"report_id\",\n",
    "    \"chunk\",\n",
    "    fc.semantic.extract(\"chunk\", enhanced_ner_schema).alias(\"chunk_entities\")\n",
    ")\n",
    "\n",
    "# Aggregate entities across chunks\n",
    "aggregated_entities = chunk_entities_df.group_by(\"report_id\").agg(\n",
    "    fc.collect_list(fc.col(\"chunk_entities\")).alias(\"all_chunk_entities\")\n",
    ")\n",
    "\n",
    "print(\"\\nChunked extraction completed for long documents\")\n",
    "print(f\"Total chunks processed: {chunk_entities_df.count()}\")\n",
    "\n",
    "# Show sample of aggregated chunk results\n",
    "print(\"\\nSample aggregated entities from chunks:\")\n",
    "aggregated_sample = aggregated_entities.select(\n",
    "    \"report_id\",\n",
    "    fc.array_size(fc.col(\"all_chunk_entities\")).alias(\"chunks_with_entities\")\n",
    ")\n",
    "aggregated_sample.show(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stage 4: Validation and Quality Assurance\n",
    "\n",
    "This cell focuses on validating the extracted entities by creating a unified view of all reports and their extracted entities. It specifically displays the CVE IDs identified in each report, allowing for quick inspection and quality assurance of the entity extraction process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Stage 4: Validation and quality assurance\n",
    "print(\"\\n✅ Stage 4: Validating extracted entities...\")\n",
    "\n",
    "# Create a unified view for validation\n",
    "all_entities_df = enhanced_df.select(\n",
    "    \"report_id\",\n",
    "    \"source\",\n",
    "    \"title\",\n",
    "    \"entities\"\n",
    ")\n",
    "\n",
    "# Show extracted CVEs\n",
    "print(\"Extracted CVE IDs:\")\n",
    "cve_summary = all_entities_df.select(\n",
    "    \"report_id\",\n",
    "    all_entities_df.entities.cve_ids.alias(\"extracted_cves\")\n",
    ")\n",
    "cve_summary.show(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stage 5: Analytics and Aggregation\n",
    "\n",
    "This cell performs analytics on the extracted entities by flattening and aggregating CVE IDs, software packages, and threat actors across all reports. It identifies the most frequently mentioned CVEs, most affected software, and most active threat actors. Finally, it generates a summary of key statistics, including total and unique CVEs, total and unique threat actors, and the number of reports processed, providing actionable security intelligence insights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Stage 5: Analytics and aggregation\n",
    "print(\"\\n📊 Stage 5: Entity analytics and insights...\")\n",
    "\n",
    "# Flatten entities for analysis\n",
    "flattened_cves = all_entities_df.select(\n",
    "    all_entities_df.entities.cve_ids.alias(\"cve_id\")\n",
    ").explode(\"cve_id\").filter(fc.col(\"cve_id\").is_not_null())\n",
    "\n",
    "flattened_software = all_entities_df.select(\n",
    "    all_entities_df.entities.software_packages.alias(\"software\")\n",
    ").explode(\"software\").filter(fc.col(\"software\").is_not_null())\n",
    "\n",
    "flattened_threats = all_entities_df.select(\n",
    "    all_entities_df.entities.threat_actors.alias(\"threat_actor\")\n",
    ").explode(\"threat_actor\").filter(fc.col(\"threat_actor\").is_not_null())\n",
    "\n",
    "# Most common CVEs\n",
    "print(\"\\nTop CVEs mentioned:\")\n",
    "cve_counts = flattened_cves.group_by(\"cve_id\").agg(\n",
    "    fc.count(\"*\").alias(\"mentions\")\n",
    ").order_by(fc.col(\"mentions\").desc())\n",
    "cve_counts.show(5)\n",
    "\n",
    "# Most affected software\n",
    "print(\"\\nMost affected software:\")\n",
    "software_counts = flattened_software.group_by(\"software\").agg(\n",
    "    fc.count(\"*\").alias(\"mentions\")\n",
    ").order_by(fc.col(\"mentions\").desc())\n",
    "software_counts.show(5)\n",
    "\n",
    "# Active threat actors\n",
    "print(\"\\nActive threat actors:\")\n",
    "threat_counts = flattened_threats.group_by(\"threat_actor\").agg(\n",
    "    fc.count(\"*\").alias(\"reports\")\n",
    ").order_by(fc.col(\"reports\").desc())\n",
    "threat_counts.show(5)\n",
    "\n",
    "# Create final comprehensive report\n",
    "print(\"\\n📋 Final Security Intelligence Summary:\")\n",
    "print(\"=\" * 70)\n",
    "\n",
    "# Summary statistics\n",
    "total_cves = flattened_cves.count()\n",
    "unique_cves = flattened_cves.select(\"cve_id\").drop_duplicates().count()\n",
    "total_threats = flattened_threats.count()\n",
    "unique_threats = flattened_threats.select(\"threat_actor\").drop_duplicates().count()\n",
    "\n",
    "print(f\"Total CVEs extracted: {total_cves} ({unique_cves} unique)\")\n",
    "print(f\"Total threat actors identified: {total_threats} ({unique_threats} unique)\")\n",
    "print(f\"Reports processed: {reports_df.count()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Actionable Intelligence: Automated Risk Assessment\n",
    "\n",
    "This cell uses a Pydantic model to define a structured risk assessment for each vulnerability report, extracting the overall risk level, recommended immediate action, and affected scope using semantic operations. \n",
    "\n",
    "It then identifies and displays high- and critical-risk vulnerabilities that require immediate attention, providing actionable security intelligence. The session is stopped at the end to clean up resources."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate actionable intelligence using semantic operations\n",
    "print(\"\\n🎯 Actionable Intelligence:\")\n",
    "\n",
    "# Define Pydantic model for risk assessment\n",
    "class ExtractedRiskInfo(BaseModel):\n",
    "    \"\"\"Directly extracted risk information from the report text. If a value is not present in the report, use an empty string.\"\"\"\n",
    "    severity_rating: str = Field(..., description=\"Explicit severity rating or risk level as stated in the report (e.g., 'critical', 'high', 'medium', 'low')\")\n",
    "    cvss_score: str = Field(..., description=\"CVSS score as stated in the report\")\n",
    "    mitigation_steps: str = Field(..., description=\"Quoted mitigation or remediation steps as stated in the report\")\n",
    "    affected_systems: str = Field(..., description=\"Exact systems, platforms, or users mentioned as affected in the report\")\n",
    "\n",
    "# Assess risk for each report\n",
    "risk_assessment_df = enhanced_df.select(\n",
    "    \"report_id\",\n",
    "    \"title\",\n",
    "    fc.semantic.extract(\"content\", ExtractedRiskInfo).alias(\"risk_assessment\")\n",
    ")\n",
    "\n",
    "# Show high-risk items\n",
    "high_risk_df = risk_assessment_df.select(\n",
    "    \"report_id\",\n",
    "    \"title\",\n",
    "    risk_assessment_df.risk_assessment.severity_rating.alias(\"risk_level\"),\n",
    "    risk_assessment_df.risk_assessment.mitigation_steps.alias(\"immediate_action\"),\n",
    "    risk_assessment_df.risk_assessment.affected_systems.alias(\"affected_scope\")\n",
    ").filter(\n",
    "    (fc.col(\"risk_level\") == \"critical\") | (fc.col(\"risk_level\") == \"high\")\n",
    ")\n",
    "\n",
    "print(\"\\nHigh-Risk Vulnerabilities Requiring Immediate Action:\")\n",
    "high_risk_df.show()\n",
    "\n",
    "# Clean up\n",
    "session.stop()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
